import paddle
import paddle.nn as nn
import numpy as np


class Generator(nn.Layer):
    '''
        生成器
    '''

    def __init__(self, n_node, node_emd_init):
        super(Generator, self).__init__()
        self.n_node = n_node
        self.node_emd_init = paddle.to_tensor(node_emd_init)
        self.embedding_matrix =paddle.create_parameter(np.array(node_emd_init).shape,dtype='float32', default_initializer=paddle.nn.initializer.Assign(node_emd_init))
        self.bias_vector = paddle.create_parameter([n_node], dtype='float32', default_initializer=paddle.nn.initializer.Assign(paddle.zeros([n_node], dtype="float32")))


    def forward(self, node_id, node_neighbor_id):
        # write down loss, optimizers in graph_gan_torch.py
        node_embedding = self.embedding_matrix[node_id]
        node_neighbor_embedding = self.embedding_matrix[node_neighbor_id]
        bias = self.bias_vector[node_neighbor_id]
        score = paddle.sum(node_embedding * node_neighbor_embedding, axis=1) + bias
        prob = nn.functional.sigmoid(score)
        prob = paddle.clip(prob, 1e-5, 1)

        return node_embedding, node_neighbor_embedding, prob

    def get_all_score(self):
        all_score = paddle.matmul(self.embedding_matrix, self.embedding_matrix.t()) + self.bias_vector

        return all_score
